import numpy as np
from chaco.api import *
from traits.api import *
from traitsui.api import *
from chaco.tools.api import LineInspector, PanTool, ZoomTool
from enable.api import ComponentEditor


class Form(HasTraits):
    N = Range(1, 100, 3)
    T = Range(2, 20.)
    tau = Range(1, 20.)
    plot = Instance(Plot)
    waveType = Enum('脉冲', '锯齿', '三角','整流')

    def traits_init(self):
        self.plot = Plot(ArrayPlotData())
        x = np.arange(-10, 10, .01)
        y = np.full_like(x, self.tau / self.T)
        self.plot.data.update_data(
            x=x,
            y=y,
        )
        self.plot.plot(("x", "y"), color='auto')
        self.update_data()

    @observe('N,tau,T,waveType')
    def update_data(self, event=None):
        x = self.plot.data.get_data('x')
        y = np.full_like(x, 0)
        match self.waveType:
            case '脉冲':
                for n in range(1, self.N + 1):
                    y += 2 / np.pi / n * np.sin(np.pi * n * self.tau / self.T) * np.cos(2 * np.pi * n / self.T * x)
            case '锯齿':
                for n in range(1, self.N + 1):
                    y += -2 * (-1.) ** n / n / np.pi * np.sin(2 * np.pi * n / self.T * x)
            case '三角':
                for n in range(1, self.N + 1):
                    y +=4 / (2*n-1)**2 / np.pi ** 2 * np.cos(2 * np.pi *(2* n-1) / self.T * x)
            case '整流':
                for n in range(1, self.N + 1):
                    y += -2 / ((2 * n) ** 2 - 1) / np.pi * np.cos(2 * np.pi *2* n / self.T * x)
        self.plot.data.update_data(
            y=y,
        )

    traits_view = View(
        UItem(
            "plot",
            editor=ComponentEditor(),
            resizable=True,
        ),
        Item('N', label='项数'),
        Item('T', label='周期'),
        Item('tau', label='宽度',visible_when='waveType=="脉冲"'),
        HGroup(
            Item('waveType', label='宽度', style='custom', editor_args={'cols':20}),
            spring,
        ),
        width=1400,
        height=1000,
        resizable=True,
    )


if __name__ == '__main__':
    form = Form()
    form.configure_traits()
